-- |
-- Module      : Stream.Generate
-- Copyright   : (c) 2018 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}

{-# OPTIONS_GHC -Wno-orphans #-}

module Stream.Generate (benchmarks) where

import Control.DeepSeq (NFData(..))
import Control.Monad.IO.Class (MonadIO)
import Data.Functor.Identity (Identity(..))
import Streamly.Internal.Data.Stream (Stream)
import Streamly.Internal.Data.Time.Units (AbsTime)

import qualified GHC.Exts as GHC
import qualified Streamly.Internal.Data.Fold as Fold
import qualified Streamly.Internal.Data.Stream as Stream

import Test.Tasty.Bench
import Stream.Common
import Streamly.Benchmark.Common
import qualified Prelude

import Prelude hiding (repeat, replicate, iterate)

-------------------------------------------------------------------------------
-- fromList
-------------------------------------------------------------------------------

{-# INLINE sourceFromList #-}
sourceFromList :: Monad m => Int -> Int -> Stream m Int
sourceFromList value n = Stream.fromList [n..n+value]

{-# INLINE sourceFromListM #-}
sourceFromListM :: MonadAsync m => Int -> Int -> Stream m Int
sourceFromListM value n = fromListM (fmap return [n..n+value])

{-# INLINE sourceIsList #-}
sourceIsList :: Int -> Int -> Stream Identity Int
sourceIsList value n = GHC.fromList [n..n+value]

{-# INLINE sourceIsString #-}
sourceIsString :: Int -> Int -> Stream Identity Char
sourceIsString value n = GHC.fromString (Prelude.replicate (n + value) 'a')

{-# INLINE readInstance #-}
readInstance :: String -> Stream Identity Int
readInstance str =
    let r = reads str
    in case r of
        [(x,"")] -> x
        _ -> error "readInstance: no parse"

-- For comparisons
{-# INLINE readInstanceList #-}
readInstanceList :: String -> [Int]
readInstanceList str =
    let r = reads str
    in case r of
        [(x,"")] -> x
        _ -> error "readInstance: no parse"

{-# INLINE replicate #-}
replicate :: Monad m => Int -> Int -> Stream m Int
replicate = Stream.replicate

-------------------------------------------------------------------------------
-- enumerate
-------------------------------------------------------------------------------

{-# INLINE sourceIntFromTo #-}
sourceIntFromTo :: Monad m => Int -> Int -> Stream m Int
sourceIntFromTo value n = Stream.enumerateFromTo n (n + value)

{-# INLINE sourceIntFromThenTo #-}
sourceIntFromThenTo :: Monad m => Int -> Int -> Stream m Int
sourceIntFromThenTo value n = Stream.enumerateFromThenTo n (n + 1) (n + value)

{-# INLINE sourceFracFromTo #-}
sourceFracFromTo :: Monad m => Int -> Int -> Stream m Double
sourceFracFromTo value n =
    Stream.enumerateFromTo (fromIntegral n) (fromIntegral (n + value))

{-# INLINE sourceFracFromThenTo #-}
sourceFracFromThenTo :: Monad m => Int -> Int -> Stream m Double
sourceFracFromThenTo value n = Stream.enumerateFromThenTo (fromIntegral n)
    (fromIntegral n + 1.0001) (fromIntegral (n + value))

{-# INLINE sourceIntegerFromStep #-}
sourceIntegerFromStep :: Monad m => Int -> Int -> Stream m Integer
sourceIntegerFromStep value n =
    Stream.take value $ Stream.enumerateFromThen (fromIntegral n) (fromIntegral n + 1)

{-# INLINE enumerateFrom #-}
enumerateFrom :: Monad m => Int -> Int -> Stream m Int
enumerateFrom count = Stream.take count . Stream.enumerateFrom

{-# INLINE enumerateFromTo #-}
enumerateFromTo :: Monad m => Int -> Int -> Stream m Int
enumerateFromTo = sourceIntFromTo

{-# INLINE enumerateFromThen #-}
enumerateFromThen :: Monad m => Int -> Int -> Stream m Int
enumerateFromThen count inp = Stream.take count $ Stream.enumerateFromThen inp (inp + 1)

{-# INLINE enumerateFromThenTo #-}
enumerateFromThenTo :: Monad m => Int -> Int -> Stream m Int
enumerateFromThenTo = sourceIntFromThenTo

-- n ~ 1
{-# INLINE enumerate #-}
enumerate :: Monad m => Int -> Int -> Stream m Int
enumerate count n = Stream.take (count + n) Stream.enumerate

-- n ~ 1
{-# INLINE enumerateTo #-}
enumerateTo :: Monad m => Int -> Int -> Stream m Int
enumerateTo count n = Stream.enumerateTo (minBound + count + n)

{-# INLINE iterate #-}
iterate :: Monad m => Int -> Int -> Stream m Int
iterate count = Stream.take count . Stream.iterate (+1)

{-# INLINE iterateM #-}
iterateM :: MonadAsync m => Int -> Int -> Stream m Int
iterateM count = Stream.take count . Stream.iterateM (return . (+1)) . return

{-# INLINE repeatM #-}
repeatM :: MonadAsync m => Int -> Int -> Stream m Int
repeatM count = Stream.take count . Stream.repeatM . return

{-# INLINE replicateM #-}
replicateM :: MonadAsync m => Int -> Int -> Stream m Int
replicateM count = Stream.replicateM count . return

{-# INLINE fromIndices #-}
fromIndices :: Monad m => Int -> Int -> Stream m Int
fromIndices value n = Stream.take value $ Stream.fromIndices (+ n)

{-# INLINE fromIndicesM #-}
fromIndicesM :: Monad m => Int -> Int -> Stream m Int
fromIndicesM value n = Stream.take value $ Stream.fromIndicesM (return <$> (+ n))

{-# INLINE _absTimes #-}
_absTimes :: MonadIO m => Int -> Int -> Stream m AbsTime
_absTimes value _ = Stream.take value Stream.absTimes

o_1_space_generation :: Int -> [Benchmark]
o_1_space_generation value =
    [ bgroup "generation"
        [ benchIOSrc "unfoldr" (sourceUnfoldr value)
        , benchIOSrc "unfoldrM" (sourceUnfoldrM value)
        , benchIOSrc "repeat" (repeat value)
        , benchIOSrc "replicate" (replicate value)
        , benchIOSrc "iterate" (iterate value)
        , benchIOSrc "iterateM" (iterateM value)
        , benchIOSrc "intFromTo" (sourceIntFromTo value)
        , benchIOSrc "intFromThenTo" (sourceIntFromThenTo value)
        , benchIOSrc "integerFromStep" (sourceIntegerFromStep value)
        , benchIOSrc "fracFromThenTo" (sourceFracFromThenTo value)
        , benchIOSrc "fracFromTo" (sourceFracFromTo value)
        , benchIOSrc "fromList" (sourceFromList value)
        , benchIOSrc "fromListM" (sourceFromListM value)
        , benchPureSrc "IsList.fromList" (sourceIsList value)
        , benchPureSrc "IsString.fromString" (sourceIsString value)
        , benchIOSrc "enumerateFrom" (enumerateFrom value)
        , benchIOSrc "enumerateFromTo" (enumerateFromTo value)
        , benchIOSrc "enumerateFromThen" (enumerateFromThen value)
        , benchIOSrc "enumerateFromThenTo" (enumerateFromThenTo value)
        , benchIOSrc "enumerate" (enumerate value)
        , benchIOSrc "enumerateTo" (enumerateTo value)
        , benchIOSrc "repeatM" (repeatM value)
        , benchIOSrc "replicateM" (replicateM value)
        , benchIOSrc "fromIndices" (fromIndices value)
        , benchIOSrc "fromIndicesM" (fromIndicesM value)

        -- fromFoldable essentially tests cons and consM which does not scale
        -- for the Stream type.
        -- , benchIOSrc "fromFoldable 16" (sourceFromFoldable 16)
        -- , benchIOSrc "fromFoldableM 16" (sourceFromFoldableM 16)
        --  XXX tasty-bench hangs benchmarking this
        -- , benchIOSrc "absTimes" $ _absTimes value
        ]
    ]

instance NFData a => NFData (Stream Identity a) where
    {-# INLINE rnf #-}
    rnf xs = runIdentity $ Stream.fold (Fold.foldl' (\_ x -> rnf x) ()) xs

o_n_heap_generation :: Int -> [Benchmark]
o_n_heap_generation value =
    [ bgroup "buffered"
    -- Buffers the output of show/read.
    -- XXX can the outputs be streaming? Can we have special read/show
    -- style type classes, readM/showM supporting streaming effects?
        [ bench "readsPrec pure streams" $
          nf readInstance (mkString value)
        , bench "readsPrec Haskell lists" $
          nf readInstanceList (mkListString value)
        ]
    ]

-------------------------------------------------------------------------------
-- Main
-------------------------------------------------------------------------------

-- In addition to gauge options, the number of elements in the stream can be
-- passed using the --stream-size option.
--
benchmarks :: String -> Int -> [Benchmark]
benchmarks moduleName size =
        [ bgroup (o_1_space_prefix moduleName) (o_1_space_generation size)
        , bgroup (o_n_heap_prefix moduleName) (o_n_heap_generation size)
        ]
